using System;
using System.Collections.Generic;
using System.Net.Sockets;
using Apache.Cassandra;

namespace SKBKontur.Cassandra.Proxy
{
    public class RetryingStorage : IStorage
    {
        private readonly IStoragePool storagePool;
        private readonly string keySpace;
        private IStorage storage;

        public RetryingStorage(IStoragePool storagePool, string keySpace)
        {
            this.storagePool = storagePool;
            this.keySpace = keySpace;
        }

        #region IStorage Members

        public void Dispose()
        {
            if (storage != null)
                storagePool.Recycle(storage, keySpace, true);
        }

        public void Insert(byte[] key, ColumnParent columnParent, ConsistencyLevel consistencyLevel,
                           params Column[] columns)
        {
            SafeCall(delegate(IStorage t)
                         {
                             t.Insert(key, columnParent, consistencyLevel, columns);
                             return true;
                         });
        }

        public Dictionary<byte[], List<ColumnOrSuperColumn>> MultigetSlice(IEnumerable<byte[]> keys,
                                                                           ColumnParent columnParent,
                                                                           SlicePredicate slicePredicate,
                                                                           ConsistencyLevel consistencyLevel)
        {
            return SafeCall(t => t.MultigetSlice(keys, columnParent, slicePredicate, consistencyLevel));
        }

        public List<ColumnOrSuperColumn> GetSlice(byte[] key, ColumnParent columnParent, SlicePredicate slicePredicate,
                                                  ConsistencyLevel consistencyLevel)
        {
            return SafeCall(t => t.GetSlice(key, columnParent, slicePredicate, consistencyLevel));
        }

        public List<KeySlice> GetIndexedSlice(ColumnParent columnParent, IndexClause indexClause, SlicePredicate column_predicate, ConsistencyLevel consistencyLevel)
        {
            return SafeCall(t => t.GetIndexedSlice(columnParent, indexClause, column_predicate, consistencyLevel));
        }

        #endregion

        private T SafeCall<T>(Func<IStorage, T> function)
        {
            if (storage == null)
                storage = storagePool.ReplaceStorage(keySpace);
            try
            {
                return function(storage);
            }
            catch (Exception)
            {
                storagePool.Recycle(storage, keySpace, false);
                storage = null;
                throw;
            }
        }
    }
}